import torch
import torchvision
import torch.utils.data as Data
# 选择MNIST的测试集用于模型的测试部分，导入方法和之前导入训练集是一样的，唯一的不同就是要把train设为False
test_data = torchvision.datasets.MNIST(root='./data/', train=False, transform=torchvision.transforms.ToTensor(), download=False)
# 将测试数据喂给模型
test_loader = Data.DataLoader(test_data, batch_size=1, shuffle=False)
# 定义使用的设备
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# 引入训练好的模型
net = torch.load('./LeNet.pkl',map_location=torch.device(device))
# 将我们的模型传至相应的设备
net.to(device)
# 在测试阶段，我们不需要对模型的参数进行更新，因此我们可以关闭自动求导功能
torch.set_grad_enabled(False)
# 使用net.eval()方法屏蔽Dropout层、冻结BN层的参数，防止在测试阶段BN层发生参数更新
net.eval()
# 获取测试集的大小
length = test_data.data.size(0)
# 定义准确率
acc = 0.0

for i, data in enumerate(test_loader):
    x, y = data
    y_pred = net(x.to(device, torch.float))
    pred = y_pred.argmax(dim=1)
    acc += (pred.data.cpu() == y.data).sum()
    print('Predict:', int(pred.data.cpu()), '|Ground Truth:', int(y.data))
acc = (acc / length) * 100
print('Accuracy: %.2f' %acc, '%')